import os
import xml.etree.ElementTree as ET
from PIL import Image

def generate_xml(image_path, xml_path):
    """生成空白的 VOC XML 标注文件"""
    # 获取图片宽高信息
    img = Image.open(image_path)
    width, height = img.size
    
    xml_path = open(xml_path, 'w')
    # 构建 XML 文件结构
    xml_path.write('<annotation>\n')
    xml_path.write('\t<folder>simple</folder>\n')
    xml_path.write('\t<filename>' + str(os.path.basename(image_path)) + '</filename>\n')
    xml_path.write('\t<size>\n')
    xml_path.write('\t\t<width>' + str(width) + ' </width>\n')
    xml_path.write('\t\t<height>' + str(height) + '</height>\n')
    xml_path.write('\t\t<depth>' + str(3) + '</depth>\n')
    xml_path.write('\t</size>\n')
    xml_path.write('</annotation>')


def main(image_folder, annotation_folder):
    """在没有对应 XML 文件的情况下生成空白的 VOC XML 标注文件"""
    # 遍历图片文件夹中的所有图片
    for image_file in os.listdir(image_folder):
        if image_file.endswith((".jpg", ".jpeg", ".png")):
            image_path = os.path.join(image_folder, image_file)
            xml_file = os.path.splitext(image_file)[0] + ".xml"
            xml_path = os.path.join(annotation_folder, xml_file)
            # 如果对应 XML 文件不存在，则生成空白的 VOC XML 标注文件
            if not os.path.isfile(xml_path):
                generate_xml(image_path, xml_path)

if __name__ == '__main__':
    # 指定图片文件夹和 XML 文件夹路径
    path = "20230507_3_50"
    for fol in os.listdir(path):
        main(path+'/'+fol, path+'/'+fol)
